import numpy as np

def calculate_energy(signal):
    return np.sum(signal ** 2)

def estimate_delay(signal, sampling_rate):
    # 增加一个显著的时延
    delayed_signal = np.roll(signal, -100)  # 向后滚动100个样本点
    correlation = np.correlate(signal, delayed_signal, mode='full')
    max_corr_index = np.argmax(correlation)
    delay = (max_corr_index - len(signal) + 1) / sampling_rate
    return delay

def beamforming(signals, delays, sampling_rate):
    max_delay = max(delays)
    adjusted_signals = []
    for signal, delay in zip(signals, delays):
        adjusted_signal = np.roll(signal, int(-delay * sampling_rate))
        adjusted_signals.append(adjusted_signal)
    beamformed_signal = np.sum(adjusted_signals, axis=0)
    return beamformed_signal

def spherical_interpolation(delays, coordinates):
    # 检查delays中的零值，替换为一个小的正数避免除零错误
    delays = np.array(delays)
    zero_delays = delays == 0
    if np.any(zero_delays):
        delays[zero_delays] = 0.0001  # 小正数替代

    weights = 1 / delays  # 延迟的倒数作为权重
    estimated_position = np.average(coordinates, axis=0, weights=weights)
    return estimated_position

def process_signal(signal, sampling_rate, methods):
    energy = calculate_energy(signal)
    delay = estimate_delay(signal, sampling_rate)
    print("能量:", energy)
    print("时延:", delay)

    signals = [signal]  # 默认信号列表
    delays = [0.0001]  # 默认延迟列表，避免除零错误

    beamformed_signal = None
    if 'beamforming' in methods:
        additional_signal = np.roll(signal, int(-delay * sampling_rate))  # 添加实际时延
        signals.append(additional_signal)
        delays.append(delay if delay != 0 else 0.0001)  # 避免延迟为零的情况
        beamformed_signal = beamforming(signals, delays, sampling_rate)
        print("波束成形结果:", beamformed_signal[:10])

    if 'spherical' in methods:
        # 确保球形插值使用波束成形后的信号
        if beamformed_signal is not None:
            # 使用波束成形后的信号更新时延和坐标
            updated_signal = beamformed_signal
            updated_delays = [estimate_delay(updated_signal, sampling_rate)]  # 重新计算时延
            coordinates = [(0, 0) for _ in updated_delays]  # 创建与更新后延迟数量匹配的坐标
            position = spherical_interpolation(updated_delays, coordinates)
            print("球形插值定位:", position)
        else:
            print("球形插值需要波束成形结果。")

# 定义声源位置和声音信号
true_position = (5, 5)  # 假设真实位置为(5, 5)
signal_length = 1000
sampling_rate = 1000
true_delay = 0.15  # 假设真实时延为0.1秒
noise_level = 0.15  # 添加的噪声水平

# 生成信号
signal = np.random.normal(0, 1, signal_length) + noise_level * np.random.normal(0, 1, signal_length)
delayed_signal = np.roll(signal, int(-true_delay * sampling_rate))  # 模拟真实时延

# 处理信号
process_signal(signal, sampling_rate, methods=[])  # 只有能量和时延
process_signal(signal, sampling_rate, methods=['beamforming'])  # 加上波束成形
process_signal(signal, sampling_rate, methods=['spherical'])  # 加上球形插值
process_signal(signal, sampling_rate, methods=['beamforming', 'spherical'])  # 所有技术